import os
import sys
import numpy as np
from typing import Optional
from pathlib import Path
import pandas as pd

script_dir = os.path.dirname(os.path.realpath(__file__))
parent_dir = os.path.dirname(os.path.dirname(script_dir))
sys.path.append(parent_dir)
import wandb
from torch.utils.data import DataLoader, random_split
from DeepPPG.dataset.dataset import ECGDataset
from DeepPPG.config import add_to_config, seed_everything
from DeepPPG.config import ResNetConfig
from DeepPPG.model.resnet1d import ResNet1d
from DeepPPG.training.seq2one_trainer import Seq2OneTrainer


def main(config_type, 
         sweep: bool, 
         log_dir: Optional[str] = None,
         model_save_dir: Optional[str] = None, 
         run_name: Optional[str] = None
         ):
    # Load config
    if isinstance(config_type, str):
        config = eval(config_type)()
    elif isinstance(config_type, (dict, ResNetConfig)):
        config = config_type
    else:
        raise ValueError(f"Invalid config type: {type(config_type)}")
    config.sweep = sweep
    # Initialize W&B
    run = wandb.init(project=config.project_name,
                     name=run_name,
                     config=None if config.sweep else dict(config), )

    print(sweep)
    if config.sweep:
        config = add_to_config(config, wandb.config)
        config.log_dir = f"{config.training_type}_{config.model_type}_log"
    else:
        config.model_save_dir = os.path.join(config.log_dir, run.id, Path("models"))

    # overwrite log_dir and model_save_dir if provided
    if log_dir is not None:
        config.log_dir = log_dir
        print("overwriting log_dir to", config.log_dir)
    if model_save_dir is not None:
        config.model_save_dir = model_save_dir
        print("overwriting model_save_dir to", config.model_save_dir)
    os.makedirs(config.log_dir, exist_ok=True)
    os.makedirs(config.model_save_dir, exist_ok=True)
    print("saving to", config.log_dir)
    seed_everything(config.seed)
    run.config.update(dict(config), allow_val_change=True)

    # Load data
    train_df = pd.read_csv(config.train_csv)
    val_df = pd.read_csv(config.val_csv)
    test_df = pd.read_csv(config.test_csv)

    train_dataset = ECGDataset(df=train_df, directory=config.data_dir, output_image=False, sig_arr_length=config.sig_length, 
                               specify_leads=config.lead_type, bandpass_lower=config.bandpass_lower, bandpass_higher=config.bandpass_higher, median_filter_size=config.median_filter_size,)
    val_dataset = ECGDataset(df=val_df, directory=config.data_dir, output_image=False, sig_arr_length=config.sig_length,
                             specify_leads=config.lead_type, bandpass_lower=config.bandpass_lower, bandpass_higher=config.bandpass_higher, median_filter_size=config.median_filter_size,)
    test_dataset = ECGDataset(df=test_df, directory=config.data_dir, output_image=False, sig_arr_length=config.sig_length,
                              specify_leads=config.lead_type, bandpass_lower=config.bandpass_lower, bandpass_higher=config.bandpass_higher, median_filter_size=config.median_filter_size,)
    
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=config.eval_batch_size, shuffle=False)
    test_loader = DataLoader(test_dataset, batch_size=config.eval_batch_size, shuffle=False)
    
    print("training set size:", len(train_dataset))
    print("validation set size:", len(val_dataset))
    print("test set size:", len(test_dataset))

    # Load model
    if config.model_type == "ResNet":
        model = ResNet1d(nOUT=config.num_classes, in_ch=config.input_channels, out_ch=config.output_channels, mid_ch=config.middle_channels)
    
    # Move model to device
    device = config.device
    model.to(device)

    # Training and validation
    trainer = Seq2OneTrainer(model=model, config=config)
    best_model, best_val_score = trainer.train(train_loader, val_loader)

    # Testing
    print("start to test ...")
    test_score = trainer.test_all(best_model, test_loader, prefix="test")

    # Log test results
    wandb.log(test_score)
    wandb.finish()

    return test_score


if __name__ == "__main__":
    import argparse
    # multiprocessing.set_start_method('spawn', force=True)

    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='ResNetConfig',
                        choices=['ResNetConfig',])
    parser.add_argument('--sweep', default=False, action='store_true')
    args = parser.parse_args()
    print("sweep : ", args.sweep)
    main(args.config, sweep=args.sweep)